import torch
import torch.nn as nn
import torch.nn.functional as F

from nets.vgg import VGG16


def get_img_output_length(width, height):
    def get_output_length(input_length):
        # input_length += 6
        filter_sizes = [2, 2, 2, 2, 2]
        padding = [0, 0, 0, 0, 0]
        stride = 2
        for i in range(5):
            input_length = (input_length + 2 * padding[i] - filter_sizes[i]) // stride + 1
        return input_length

    return get_output_length(width) * get_output_length(height)


class Siamese(nn.Module):
    def __init__(self, input_shape, pretrained=False):
        super(Siamese, self).__init__()
        self.vgg = VGG16(pretrained, input_shape[-1])
        del self.vgg.avgpool
        del self.vgg.classifier

        flat_shape = 512 * get_img_output_length(input_shape[1], input_shape[0])
        print(flat_shape)
        self.fully_connect1 = torch.nn.Linear(25088, 512)
        self.fully_connect2 = torch.nn.Linear(512, 1)

        # 空间变换器定位 - 网络
        self.localization = nn.Sequential(
            nn.Conv2d(3, 8, kernel_size=5, padding=1, stride=2),
            nn.MaxPool2d(2, stride=2),
            nn.ReLU(True),
            nn.Conv2d(8, 16, kernel_size=5, padding=1, stride=2),
            nn.MaxPool2d(2, stride=2),
            nn.ReLU(True),
            nn.Conv2d(16, 32, kernel_size=3, stride=1),
            nn.MaxPool2d(2, stride=2),
            nn.ReLU(True),
            nn.Conv2d(32, 32, kernel_size=2, stride=1),
            nn.MaxPool2d(2, stride=1),
            nn.ReLU(True)
        )
        # 3 * 2 affine矩阵的回归量
        self.fc_loc = nn.Sequential(
            nn.Linear(32 * 3 * 3, 32),
            nn.ReLU(True),
            nn.Linear(32, 3 * 2)
        )
        # 使用身份转换初始化权重/偏差
        self.fc_loc[2].weight.data.zero_()
        self.fc_loc[2].bias.data.copy_(torch.tensor([1, 0, 0, 0, 1, 0], dtype=torch.float))  # 网络层参数初始化，默认图片保持原样

    # 空间变换器网络转发功能
    def stn(self, x):
        xs = self.localization(x)
        xs = xs.view(-1, 32 * 3 * 3)  # -1为占位
        theta = self.fc_loc(xs)
        theta = theta.view(-1, 2, 3)
        grid = F.affine_grid(theta, x.size())  # 计算初步调整坐标
        x = F.grid_sample(x, grid)  # 根据自己定义的填充规则（一般用双线性插值）来填充
        return x

    def forward(self, x):
        x1, x2 = x
        x1 = self.stn(x1)
        x2 = self.stn(x2)
        # ------------------------------------------#
        #   我们将两个输入传入到主干特征提取网络
        # 这里x1为left x2为right
        # ------------------------------------------#
        x1 = self.vgg.features(x1)  # 这里只使用了主干网络
        x2 = self.vgg.features(x2)
        # -------------------------#
        #   相减取绝对值
        # -------------------------#
        x1 = torch.flatten(x1, 1)
        x2 = torch.flatten(x2, 1)
        x = torch.abs(x1 - x2)
        # -------------------------#
        #   进行两次全连接
        # -------------------------#
        x = self.fully_connect1(x)
        x = self.fully_connect2(x)
        return x


class ContrastiveLoss(torch.nn.Module):

    def __init__(self, margin):
        super(ContrastiveLoss, self).__init__()
        self.margin = margin

    def forward(self, x1, x2, y):
        dist = F.pairwise_distance(x1, x2)
        total_loss = (1 - y) * torch.pow(dist, 2) + \
                     y * torch.pow(torch.clamp_min_(self.margin - dist, 0), 2)
        loss = torch.mean(total_loss)
        return loss


if __name__ == "__main__":


    ContrastiveLoss = ContrastiveLoss(0.3)
    a = torch.tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3]])
    b = torch.tensor([[1, 12, 3], [1, 20, 30], [10, 20, 3]])
    y = 1#距离
    ContrastiveLoss(a,b,y)
    model = Siamese([448, 448, 3], True)

    # from torchsummary import summary
    # summary(model, input_size=(2,3, 448,448), batch_size=16)

    # fms = model(torch.randn( 3, 448,448))
    #
    from ptflops import get_model_complexity_info

    flops, params = get_model_complexity_info(model, (3, 448, 448), as_strings=True,
                                              print_per_layer_stat=True, batch_size=2)  # 不用写batch_size大小，默认batch_size=1
    print('Flops:  ' + flops)
    print('Params: ' + params)